from gauss import gauss
import numpy as np
class backChargeGauss(gauss):
  def setKey(self,charge,method='hf',basis='3-21g'):
    self.key = '#p force nosymm charge ' + method + '/' + basis + '\n\n'
    back = []
    backCrd = charge[0]
    backCharge = charge[1]
    for i in xrange(len(backCharge)):
      back.append(' {0:14.8f}{1:14.8f}{2:14.8f}{3:14.8f}\n'.format(backCrd[i][0],backCrd[i][1],\
                 backCrd[i][2],backCharge[i]))
    back[-1] += '\n'
    self.others = ''.join(back)

  def read(self,name):
    log = open(name,'r')
    records = []
    while 1:
      line = log.readline()
      if 'Self energy of the charges' in line:
        self.CalcChargeEnergy = float(line[30:50])
      elif 'SCF Done:' in line:
        self.CalcEnergy = float(line[23:39])
      elif 'Forces (Hartrees/Bohr)' in line:
        log.readline();log.readline()
        while 1:
          line = log.readline()
          if '---' in line:
            break
          records.append(line[24:])
      if line == '':
        break
    records = ''.join(records)
    forces = np.fromstring(records, dtype=float, sep=' ')
    self.CalcForce = forces.reshape(forces.size/3,3)


if __name__ == '__main__':
    spin = 1 ; charge = 1
    name = 'bbb'
    atom = [[1,1,1],[2,2,2],[3,3,3]]
    atomName = ['H','O','H']
    cube = (atom,atomName)
    back = [[4,4,4],[5,5,5],[6,6,6]]
    backCharge = [1,1,1]
    backgroud = (back,backCharge)
    a = backChargeGauss()
    a.setBasic(name,spin,charge,cube)
    a.setKey(backgroud)
    a.write()
    a.read('zn_p_001002.log')
    print a.CalcChargeEnergy,a.CalcEnergy
    print a.CalcForce
